from langchain_core.tools import tool
from langchain.chat_models import init_chat_model
import os

key = os.getenv("OPENAI_API_KEY")

llm = init_chat_model(
    model="gpt-4o-mini",
    model_provider="openai",
    base_url="https://api.zetatechs.com/v1",
    api_key=key
)

@tool
def add(a: int, b: int) -> int:
    """Adds a and b."""
    return a + b


@tool
def multiply(a: int, b: int) -> int:
    """Multiplies a and b."""
    return a * b


tools = [add, multiply]
# llm_forced_to_multiply = llm.bind_tools(tools, tool_choice="multiply")
llm_forced_to_multiply = llm.bind_tools(tools)
print(llm_forced_to_multiply.invoke("what is 2 + 4"))